Deterministic MoE Router Kernel By Removing Global Atomic Ops by zhongbozhu · Pull Request #3587 · NVIDIA/TransformerEngine

zhongbozhu · 2026-09-29T18:25:18Z

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes NVIDIA/Megatron-LM#7221

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

Change A
Change B

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

greptile-apps · 2026-09-29T18:31:35Z

Confidence Score: 5/5

[Medium impact] The reviewed changes appear safe to merge; no blocking issue was found.

Summary

Adds an optional fixed-order reduction for MoE auxiliary loss while keeping the existing atomic path as the default.

MoE auxiliary loss can use a fixed-order reduction.
The fused MLA module keeps its imports in source order.

No actionable new issues were found. Tests were not run. No previous review threads were supplied, and their history could not be retrieved.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[fused_moe_aux_loss] --> B{deterministic?}
  B -->|No| C[Existing atomic reduction]
  B -->|Yes| D[Get workspace]
  D --> E[Each block writes one partial sum]
  E --> F[Fixed-order final sum]
  C --> G[Loss and saved backward coefficient]
  F --> G

_{Reviews (7) · Last reviewed commit: "Merge branch 'main' into fix/determinist..." · Reviewed by Greptile}

ptrendx · 2026-09-29T23:25:47Z

+ *
+ *  \return Maximum required workspace size in bytes.
+ */
+size_t nvte_get_moe_aux_loss_workspace_size();


The way we handle the workspace in the other APIs is similar to CUB - we run the APIs with the nullptr arguments for the workspace once (which populates the other fields), then run it again with the proper workspace. Please align this API to that.

ZhiyuLi-Nvidia · 2026-10-12T06:52:19Z

Thanks a lot @zhongbozhu.

The kernel is deterministic on one GPU. For multi-node runs, does determinism also hold with multi-node job, do we need the same placement and a fixed NCCL algorithm? Good to put it in the docs if there's anything worth noting.

zhongbozhu · 2026-10-15T07:30:03Z

/te-ci pytorch

ZhiyuLi-Nvidia · 2026-10-19T23:14:11Z

Verified on my end:

mcore determinism in unit tests
e2e run https://forge.coreweave.com/wandb/nvidia/determinism-zhiyul?nw=tv4seg54n2
Feel free to move forward.

zhongbozhu · 2026-10-20T00:30:42Z

/te-ci pytorch

vthumbe1503

LGTM, with minor suggestion on workspace contract

vthumbe1503 · 2026-10-26T03:29:57Z

-import transformer_engine_torch as tex
 from packaging.version import Version as PkgVersion

+import transformer_engine_torch as tex


This file change is not needed right?

that seems to be linter complaint?

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> Signed-off-by: zhongboz <zhongboz@nvidia.com>

Query workspace shape and dtype through the deterministic forward APIs, retain stream and CUDA Graph isolation, and cover metadata queries and invalid scratch buffers. Fix the stream cache-key lint warning and upstream import ordering. Validated on GB200: 201 TE tests passed (12 invalid top-k cases skipped), 6 C API tests passed, and 40 MCore checks passed with 1 expected atomic negative-control XPASS. Full PyTorch QA lint, pre-commit, and license checks passed. Signed-off-by: zhongboz <zhongboz@nvidia.com>

Signed-off-by: zhongboz <zhongboz@nvidia.com>

Apply Varun Thumbe's workspace-contract changes from #1. Require the queried 256-element FP32 capacity for every input shape and update the C API workspace tests. Retain the existing MLA import grouping required by Python lint. Original author: Varun Thumbe <vthumbe@nvidia.com> Source: 9a3e764 Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Signed-off-by: zhongboz <zhongboz@nvidia.com>

zhongbozhu · 2026-10-28T08:30:32Z

/te-ci pytorch

vthumbe1503 · 2026-10-30T05:15:02Z

/te-ci L0

vthumbe1503 · 2026-10-30T17:02:54Z

Relevant tests are passing. Pipeline: 77803646

* initial draft Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> Signed-off-by: zhongboz <zhongboz@nvidia.com> * Align deterministic MoE aux-loss workspace queries with the common API Query workspace shape and dtype through the deterministic forward APIs, retain stream and CUDA Graph isolation, and cover metadata queries and invalid scratch buffers. Fix the stream cache-key lint warning and upstream import ordering. Validated on GB200: 201 TE tests passed (12 invalid top-k cases skipped), 6 C API tests passed, and 40 MCore checks passed with 1 expected atomic negative-control XPASS. Full PyTorch QA lint, pre-commit, and license checks passed. Signed-off-by: zhongboz <zhongboz@nvidia.com> * Test deterministic MoE aux-loss integer boundaries Signed-off-by: zhongboz <zhongboz@nvidia.com> * Simplify deterministic MoE aux-loss workspace contract Apply Varun Thumbe's workspace-contract changes from zhongbozhu#1. Require the queried 256-element FP32 capacity for every input shape and update the C API workspace tests. Retain the existing MLA import grouping required by Python lint. Original author: Varun Thumbe <vthumbe@nvidia.com> Source: zhongbozhu@9a3e764 Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Signed-off-by: zhongboz <zhongboz@nvidia.com> --------- Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> Signed-off-by: zhongboz <zhongboz@nvidia.com> Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Co-authored-by: Varun Thumbe <vthumbe@nvidia.com> Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com> (cherry picked from commit fd2ab0c)

zhongbozhu requested review from ksivaman and ptrendx as code owners September 29, 2026 18:25

github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Sep 29, 2026

ptrendx reviewed Sep 29, 2026

View reviewed changes

ptrendx added the 2.21 label Sep 30, 2026

This was referenced Oct 1, 2026

[Determinism] Support deterministic fused MoE auxiliary loss NVIDIA/Megatron-LM#7697

Open

[Determinism] Add deterministic TE fused loss support NVIDIA/Megatron-LM#7221

Closed

zhongbozhu force-pushed the fix/deterministic-moe-aux-loss branch from ecdbe5d to 17c35fa Compare October 8, 2026 09:47

zhongbozhu requested a review from cyanguwa as a code owner October 8, 2026 09:47

zhongbozhu force-pushed the fix/deterministic-moe-aux-loss branch from 17c35fa to 07887f8 Compare October 15, 2026 07:29

zhongbozhu force-pushed the fix/deterministic-moe-aux-loss branch from 07887f8 to a5d4c78 Compare October 19, 2026 23:12

vthumbe1503 reviewed Oct 26, 2026

View reviewed changes

zhongbozhu added 3 commits October 27, 2026 22:49

initial draft

c1c64f2

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> Signed-off-by: zhongboz <zhongboz@nvidia.com>

Test deterministic MoE aux-loss integer boundaries

23ce45f

Signed-off-by: zhongboz <zhongboz@nvidia.com>

zhongbozhu force-pushed the fix/deterministic-moe-aux-loss branch from a5d4c78 to 23ce45f Compare October 28, 2026 07:05

vthumbe1503 approved these changes Oct 29, 2026

View reviewed changes

Merge branch 'main' into fix/deterministic-moe-aux-loss

d26dc22

vthumbe1503 merged commit fd2ab0c into NVIDIA:main Oct 30, 2026
33 of 38 checks passed

Conversation

zhongbozhu commented Sep 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented Sep 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Summary

Diagram

Uh oh!

ptrendx Sep 29, 2026

Choose a reason for hiding this comment

Uh oh!

zhongbozhu Oct 9, 2026

Choose a reason for hiding this comment

Uh oh!

ZhiyuLi-Nvidia commented Oct 12, 2026

Uh oh!

zhongbozhu commented Oct 15, 2026

Uh oh!

ZhiyuLi-Nvidia commented Oct 19, 2026

Uh oh!

zhongbozhu commented Oct 20, 2026

Uh oh!

vthumbe1503 left a comment

Choose a reason for hiding this comment

Uh oh!

vthumbe1503 Oct 26, 2026

Choose a reason for hiding this comment

Uh oh!

zhongbozhu Oct 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

zhongbozhu commented Oct 28, 2026

Uh oh!

vthumbe1503 commented Oct 30, 2026

Uh oh!

vthumbe1503 commented Oct 30, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

zhongbozhu commented Sep 29, 2026 •

edited

Loading

greptile-apps Bot commented Sep 29, 2026 •

edited

Loading

zhongbozhu Oct 28, 2026 •

edited

Loading